#include "ap_theta_algorithm.hpp"
#include <limits>
#include <queue>
#include <cmath>

int APThetaAlgorithm::heuristic(const Point &a, const Point &b)
{
  return std::abs(a.x - b.x) + std::abs(a.y - b.y);
}

double APThetaAlgorithm::angleBetween(const Point &a, const Point &b, const Point &c)
{
  double ab_x = b.x - a.x;
  double ab_y = b.y - a.y;
  double ac_x = c.x - a.x;
  double ac_y = c.y - a.y;

  double dot_product = ab_x * ac_x + ab_y * ac_y;
  double magnitude_ab = std::sqrt(ab_x * ab_x + ab_y * ab_y);
  double magnitude_ac = std::sqrt(ac_x * ac_x + ac_y * ac_y);

  double cos_theta = dot_product / (magnitude_ab * magnitude_ac);
  return std::acos(cos_theta) * (180.0 / CV_PI);
}

std::vector<Point> APThetaAlgorithm::search(std::vector<std::vector<int>> &grid, Point start, Point end, cv::Mat &image, int occ_size)
{
  int rows = grid.size();
  int cols = grid[0].size();
  std::vector<std::vector<bool>> visited(rows, std::vector<bool>(cols, false));
  std::vector<std::vector<Point>> parent(rows, std::vector<Point>(cols, {-1, -1}));
  std::vector<std::vector<int>> g_cost(rows, std::vector<int>(cols, std::numeric_limits<int>::max()));
  std::vector<std::vector<double>> lb(rows, std::vector<double>(cols, -std::numeric_limits<double>::infinity()));
  std::vector<std::vector<double>> ub(rows, std::vector<double>(cols, std::numeric_limits<double>::infinity()));

  auto cmp = [](const std::pair<Point, double> &left, const std::pair<Point, double> &right)
  { return left.second > right.second; };
  std::priority_queue<std::pair<Point, double>, std::vector<std::pair<Point, double>>, decltype(cmp)> pq(cmp);
  pq.push({start, 0});
  g_cost[start.y][start.x] = 0;

  int dx[] = {-1, 0, 1, 0, -1, 1, -1, 1};
  int dy[] = {0, 1, 0, -1, -1, -1, 1, 1};

  while (!pq.empty())
  {
    Point current = pq.top().first;
    pq.pop();

    if (visited[current.y][current.x])
      continue;
    visited[current.y][current.x] = true;

    if (current == end)
    {
      std::vector<Point> path;
      while (current.x != -1 && current.y != -1)
      {
        path.emplace_back(current);
        current = parent[current.y][current.x];
      }
      std::reverse(path.begin(), path.end());
      return path;
    }

    for (int i = 0; i < 8; ++i)
    {
      int nx = current.x + dx[i];
      int ny = current.y + dy[i];

      if (nx >= 0 && nx < cols && ny >= 0 && ny < rows && grid[ny][nx] == 0)
      {
        double new_g_cost = g_cost[current.y][current.x] + 1;
        double h_cost = heuristic({nx, ny}, end);
        double f_cost = new_g_cost + h_cost;

        if (new_g_cost < g_cost[ny][nx])
        {
          double angle = angleBetween(parent[current.y][current.x], current, {nx, ny});
          if (angle >= lb[current.y][current.x] && angle <= ub[current.y][current.x])
          {
            g_cost[ny][nx] = new_g_cost;
            pq.push({{nx, ny}, f_cost});
            parent[ny][nx] = current;

            if (!(nx == end.x && ny == end.y))
            {
              cv::Rect roi(nx * occ_size, ny * occ_size, occ_size, occ_size);
              image(roi) = cv::Vec3b(0, 255, 0);
              cv::imshow("AP Theta* Expansion", image);
              cv::waitKey(1);
            }
          }
        }
      }
    }
  }

  return {};
}